# -*- coding=utf-8 -*-
# Created Time: 2016年09月18日 星期日 09时42分53秒

"""
封装常用 excel 操作
"""

from openpyxl import Workbook
import StringIO


class Excelib(object):

    def __init__(self, file_name):
        self.file_name = file_name
        self.wb = Workbook()

    def create(self, data_list):
        """ 根据数据 data_list 创建一个 新的 excel 文件

        支持创建多个 sheet
        data_list 必须是固定的格式:
        [
            {
                'sheet':'sheet_1',                   # 第一个 sheet
                'data':[                             # 数据
                    ['表头1','表头2','表头3'],       # 表头
                    ['zoo', 'tiger', 'bird'],        # 数据项
                    [1, 100, '2016-09-18'],          # 数据项
                    ...
                ]
            },
            ...                                      # 第二个 sheet
        ]
        """

        if not data_list: return

        for i,data in enumerate(data_list):
            if i == 0:
                ws = self.wb.active
            else:
                ws = self.wb.create_sheet()

            ws.title = data.get('sheet', 'foo')
            for bar in data.get('data', []):
                ws.append(bar)

    def save2file(self):
        """ 保存数据到 Excel 文件 """

        self.wb.save(self.file_name)

    def save2io(self):
        """ 保存数据到 流 """

        sio = StringIO.StringIO()
        self.wb.save(sio)
        return sio


def main():
    """ 测试 """

    data_list = [
        {
            'sheet':u'第一个表格',
            'data':[
                ['姓名', '年龄', '体重'],
                ['god', '22', 90],
                ['libo', '99', 190]
            ]
        },
        {
            'sheet':u'第二个表格',
            'data':[
                ['姓名', '学校', '专业'],
                ['god', '电子科大', '土木工程'],
                ['libo', '电子科大成都学院', '软件测试']
            ]
        }
    ]

    excelib = Excelib('export_file.xlsx')
    excelib.create(data_list)
    excelib.save2file()

if __name__ == '__main__':
    main()
